#!/usr/bin/env python

# loading dicom datasets

import dicom
from os import listdir
from numpy import *
from scipy.io.numpyio import fread


def read_mha(filename):
    f = open(filename)
    
    # read header
    header = {}
    
    while True:
        key, value = [x.strip() for x in f.readline().split('=')]
        header[key] = value
        if key == 'ElementDataFile':
            assert value == 'LOCAL'
            break
    
    assert header['BinaryData'] == 'True'
    assert header['ElementType'] == 'MET_FLOAT'

    w, h, d = map(int, header['DimSize'].split())
    data = zeros((h, w, d), dtype = int16, order = "F")
    
    # read data
    for i in range(d):
        slice = fread(f, h*w, 'f')
        slice = reshape(slice, (h, w))
        data[:,:,i] = slice
    
    f.close()
    return data


def dicom_sort(files):
    for f in files: 
        assert f.endswith(".dcm")
    
    def n(f):
        f = f[:-4]
        i = len(f) - 1
        while i >= 0 and f[i] >= '0' and f[i] <= '9': i -= 1
        return int(f[i+1:])

    def compare(x, y):
        return cmp(n(x), n(y))
    
    return sorted(files, compare)


def read_dicom_slice(f):
    # print "read_dicom_slice:", f
    dcm = dicom.read_file(f)
    img = dcm.PixelArray
    assert img.dtype == int16
    
    if hasattr(dcm, 'RescaleSlope'):
        assert dcm.RescaleSlope == 1
        
    if hasattr(dcm, 'RescaleIntercept'):
        img += int16(dcm.RescaleIntercept)
    
    return img
    

def read_dicom_data(dcm):
    # is it an MHA cube?
    if dcm.endswith('.mha'):
        return read_mha(dcm)
        
    # is it single dicom file?
    if dcm.endswith('.dcm'):
        return read_dicom_slice(dcm) # this will be a 3d array, not a slice

    # read dicom directory
    files = dicom_sort(listdir(dcm))

    # init numpy array for the voxel data
    dcm = dicom.read_file(dcm + "/" + files[0])
    height, width = dcm.PixelArray.shape
    depth = len(files)
    data = zeros((height, width, depth), dtype = int16, order = "F")

    # load voxel data
    for i, f in enumerate(files):
        data[:,:,i] = read_dicom_slice(dcm + "/" + f)
        
    return data


def read_4d_dicom_data(dicomDir4d):
    # read dicom directory
    dirs = [d for d in listdir(dicomDir4d)]
    dirs.sort()
    print dirs
    dirs = [dicomDir4d + '/' + d for d in dirs]
    
    # are they MHA by any chance?
    if dirs[0].endswith('.mha'):
        cubes = [read_mha(f) for f in dirs]
        h, w, d = cubes[0].shape
        return dstack(cubes), (h, w, d, len(cubes))

    # init numpy array for the voxel data
    files = listdir(dirs[0])
    dcm = dicom.read_file(dirs[0] + "/" + files[0])
    height, width = dcm.PixelArray.shape
    depth = len(files)
    frames = len(dirs)
    data = zeros((height, width, depth * frames), dtype = int16, order = "F")
    print height, width, depth, frames

    # load voxel data
    i = 0
    for d in dirs:
        print "loading:", d
        files = dicom_sort(listdir(d))
        assert len(files) == depth
        for f in files:
            data[:,:,i] = read_dicom_slice(d + "/" + f)
            i += 1
            
    return data, (height, width, depth, frames)


if __name__ == '__main__':
    from sys import argv
    v = read_dicom_data(argv[1])
    print 'shape', v.shape
    x = v.ravel()
    print 'min', x.min(), 'max', x.max(), 'avg', x.sum() / len(x)
    
